# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/23 10:10 
@Author  : ZhangShenao 
@File    : vector_store_service.py 
@Desc    : 向量数据库服务
"""
from typing import List

from injector import inject
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_core.documents import Document
from langchain_core.vectorstores import VectorStoreRetriever
from langchain_pinecone import PineconeVectorStore


@inject
class VectorStoreService:
    """向量数据库服务,封装Pinecone向量数据库的操作"""

    def __init__(self):
        """初始化项链数据库服务"""

        # 初始化Embedding模型
        embeddings = ZhipuAIEmbeddings(model='embedding-2')

        # 初始化Pinecone数据库
        self.__vector_store = PineconeVectorStore(
            index_name='dataset',  # 指定索引名称
            embedding=embeddings,  # 指定Embedding模型
            namespace='llm-ops',  # 指定Namespace
            text_key='text',  # 指定文本属性key
        )

    def as_retriever(self) -> VectorStoreRetriever:
        """获取向量数据库检索器"""

        return self.__vector_store.as_retriever(search_kwargs={'k': 2})

    @classmethod
    def join_document_page_contents(cls, docs: List[Document]) -> str:
        """
        将文档列表中的每个文档的page_content属性合并为一个字符串
        多个文档之间以换行符分割
        :param docs: 文档列表
        :return: 文档内容合并后的字符串
        """

        return '\n'.join([doc.page_content for doc in docs])
